# -*- coding: utf-8 -*-
"""
******* 文档说明 ******

https://www.jianshu.com/p/0fb830696a1c

# 当前项目: MovieLens
# 创建时间: 2020-09-24  09:49
# 开发作者: Vincent
# 版    本: V1.0
"""
import os
import sys
import logging
import numpy as np
import pandas as pd
import tensorflow as tf

from ncf_model.dataset import Dataset as NCFDataset
from ncf_model.ncf_singlenode import NCF
from common.python_splitters import python_stratified_split
from common.python_evaluation import map_at_k, ndcg_at_k, precision_at_k, recall_at_k

print("System version: {}".format(sys.version))
print("Numpy  version: {}".format(np.__version__))
print("Pandas version: {}".format(pd.__version__))
print("Tensorflow  version: {}".format(tf.__version__))

logging.basicConfig(level=logging.DEBUG,
                    format='%(asctime)s %(levelname)-8s %(message)s')

# top k items to recommend
TOP_K = 10
EPOCHS = 50
BATCH_SIZE = 256
SEED = 42


# 读取数据
def get_data(data_path):
    ratings_title = ['userID','itemID', 'rating', 'timestamp']
    data = pd.read_table(os.path.join(data_path, 'ratings.dat'), sep='::',
                            header=None, names=ratings_title, engine = 'python')
    data['rating'] = data['rating'].astype(np.float32)
    # print(data.head())

    movies_title = ['MovieID', 'Title', 'Genres']
    movies = pd.read_table(os.path.join(data_path, 'movies.dat'),
                           sep='::', header=None, names=movies_title, engine='python')

    return data, movies


# 打印测试结果
def print_test(y_true, y_pre, item_map, top_k):
    """
    :param y_true:
    :param y_pre:
    :param item_map:
    :param top_k:
    :return:
    """
    y_true = y_true.sort_values(by=['rating'], ascending=False).copy()

    eval_map = map_at_k(y_true, y_pre, col_user='userID', col_item='itemID', col_rating='rating', k=top_k)
    eval_ndcg = ndcg_at_k(y_true, y_pre, col_user='userID', col_item='itemID', col_rating='rating', k=top_k)
    eval_precision = precision_at_k(y_true, y_pre, col_user='userID', col_item='itemID', col_rating='rating', k=top_k)
    eval_recall = recall_at_k(y_true, y_pre, col_user='userID', col_item='itemID', col_rating='rating', k=top_k)

    print('map_at_k:        {:.3f}'.format(eval_map))
    print('ndcg_at_k:       {:.3f}'.format(eval_ndcg))
    print('precision_at_k:  {:.3f}'.format(eval_precision))
    print('recall_at_k:     {:.3f}\n'.format(eval_recall))

    # 真实值 Item 字典
    y_true_item = {v: i for i, v in enumerate(y_true['itemID'])}

    for i in range(len(y_true)):
        item_i = y_true['itemID'].iloc[i]
        rating_i = y_true['rating'].iloc[i]

        print('{:3s} {:5d}  {:.3f}  {}      {}'.format('', item_i, rating_i,
                                                       item_map.loc[[item_i], ['Title']].values[0][0],
                                                       item_map.loc[[item_i], ['Genres']].values[0][0]))

    print('- '*20, 'Predict', ' -'*20)
    for i in range(len(y_pre)):
        item_i = y_pre['itemID'].iloc[i]
        rating_i = y_pre['prediction'].iloc[i]

        top_i = str(y_true_item[item_i]) if item_i in y_true_item else ''

        print('{:3s} {:5d}  {:.3f}  {}      {}'.format(top_i, item_i, rating_i,
                                                       item_map.loc[[item_i], ['Title']].values[0][0],
                                                       item_map.loc[[item_i], ['Genres']].values[0][0]))


def main():
    # 读取数据
    data, movies_map = get_data(data_path = r'D:\Data\Recommend\ml-1m')
    data = data[:1000]

    # 拆分训练集、测试集 分组拆开
    # train, test = python_chrono_split(data, 0.75)
    train, test = python_stratified_split(data, ratio=0.75, col_user='userID', col_item='itemID', seed=42)
    print('TrainData: {}\n{}'.format(train.shape, train.head()))
    print('TestData : {}\n{}'.format(test.shape, test.head()))

    data = NCFDataset(train=train, test=test, seed=SEED)

    print(data.n_users)
    print(data.n_items)

    model = NCF (
        n_users=data.n_users,
        n_items=data.n_items,
        model_type="NeuMF",
        n_factors=4,
        layer_sizes=[16,8,4],
        n_epochs=EPOCHS,
        batch_size=BATCH_SIZE,
        learning_rate=1e-3,
        verbose=10,
        seed=SEED
    )

    # ################################ 模型训练
    import time
    start_time = time.time()
    model.fit(data)
    train_time = time.time() - start_time
    print("Took {} seconds for training.".format(train_time))

    # ################################ 模型预测
    start_time = time.time()

    # 预测结果
    y_predict = list()
    # 所有Items集合
    item_all = set(train.itemID.unique())
    for user_i in train.userID.unique():
        # 当前用户训练数据中的 Items
        train_items = set(train['itemID'][train['userID'] == user_i].values)
        # 待预测的 Items数据
        item = list(item_all - train_items)
        # 用户数据
        user = [user_i] * len(item)

        # 模型预测各item 得分
        score = list(model.predict(user, item, is_list=True))
        # 排序
        pred_user_info = sorted(zip(user, item, score), key=lambda x:x[2], reverse=True)
        # TopK
        y_predict.extend(pred_user_info[:TOP_K])

    y_predict = pd.DataFrame(y_predict, columns=['userID', 'itemID', 'prediction'])

    test_time = time.time() - start_time
    print("Took {} seconds for prediction.".format(test_time))

    eval_map = map_at_k(test, y_predict, col_user='userID', col_item='itemID', col_rating='rating', k=TOP_K)
    print(eval_map)

    # 打印测试结果
    print_test(test[test['userID'] == 7], y_predict[y_predict['userID'] == 7], movies_map, top_k=TOP_K)


    # ################################################## 模型预测
    # start_time = time.time()
    # users, items, preds = [], [], []
    # item = list(train.itemID.unique())
    # for user in train.userID.unique():
    #     user = [user] * len(item)
    #     users.extend(user)
    #     items.extend(item)
    #     preds.extend(list(model.predict(user, item, is_list=True)))
    #
    # all_predictions = pd.DataFrame(data={"userID": users, "itemID":items, "prediction":preds})
    #
    # # 清除训练数据中标签
    # merged = pd.merge(train, all_predictions, on=["userID", "itemID"], how="outer")
    # y_predict_temp = merged[merged.rating.isnull()].drop('rating', axis=1)
    #
    # y_predict_top_k = list()
    # for user_i in y_predict_temp['userID'].unique():
    #     # print(user_i)
    #     y_user_i = y_predict_temp[y_predict_temp['userID']==user_i].copy()
    #     y_user_i = y_user_i.sort_values(by=['prediction'], ascending=False)
    #     y_user_i = y_user_i[['userID','itemID', 'prediction']]
    #     y_predict_top_k.append(y_user_i[:TOP_K])
    #
    # y_predict = pd.concat(y_predict_top_k)
    #
    # test_time = time.time() - start_time
    # print("Took {} seconds for prediction.".format(test_time))
    #
    # eval_map = map_at_k(test, y_predict, col_user='userID', col_item='itemID', col_rating='rating', k=TOP_K)
    # print(eval_map)
    #
    # # 打印测试结果
    # print_test(test[test['userID'] == 7], y_predict[y_predict['userID'] == 7], movies_map, top_k=TOP_K)


if __name__ == '__main__':
    main()
    print(os.path.abspath(os.path.dirname(__file__)))
